#!/usr/bin/env python

import csv
import itertools
import ssh
import subprocess
import sys
import threading
import time
import Queue

# key file must be in same directory
PRIV_KEY = 'zhtemp1'
OUTPUT_FILE = 'colocations.csv'

MIN_NODES = 8 # 8
MAX_NODES = 256 # 256

# ec2 api tool commands
ec2describe = 'ec2-describe-instances'
ec2launch = 'ec2-run-instances ami-714ba518 -t t1.micro -k {0} ' \
                 '-z us-east-1a'.format(PRIV_KEY).split()

def main():
    current_ids = []
    num_nodes = MIN_NODES
    outfile = open(OUTPUT_FILE, 'wb')
    csv_writer = csv.writer(outfile, delimiter=',', quoting=csv.QUOTE_NONE)
    write_queue = Queue.Queue()
    threads = []

    print "launching %d nodes" % num_nodes
    # launch MIN_NODES instances
    for i in range(0, num_nodes):
        plaunch = subprocess.Popen(ec2launch)
        plaunch.wait()

    while num_nodes <= MAX_NODES:
        csv_writer.writerow([num_nodes])
        print "waiting for instances to launch"
        instance_info = wait_for_launch() # saves a call to ec2describe

        # setup the instances
        print "setting up the instances"
        instance_ids = get_instance_ids(instance_info) # all ids
        setup_ids = [id for id in instance_ids if id not in current_ids]
        threads = [threading.Thread(target=setup_instance, args=(id,)) \
                        for id in setup_ids]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        current_ids.extend(setup_ids)

        # collect the stats
        print "running bandwidth tests"
        pairs = [p for p in itertools.combinations(current_ids, 2)]
        test_pairs = []
        used = []
        threads[:] = []

        while pairs:
            for p in pairs:
                if p[0] in used or p[1] in used:
                    continue
                test_pairs.append(p)
                map(used.append, p)

            threads = [threading.Thread(target=run_test, args=(tp, write_queue)) \
                            for tp in test_pairs]
            for t in threads:
                t.start()
            for t in threads:
                t.join()

            for tp in test_pairs:
                try:
                    pairs.remove(tp)
                except ValueError:
                    print "%s not in pairs" % (tp,)
                    sys.exit(1)

            # clear the lists
            used[:] = []
            test_pairs[:] = []
            threads[:] = []
            assert len(used) == len(test_pairs) == len(threads) == 0

            while not write_queue.empty():
                row = write_queue.get()
                csv_writer.writerow(row)

        if num_nodes >= MAX_NODES:
            break

        # launch NUM_NODES more instances
        print "launching %d nodes" % num_nodes
        for i in range(0, num_nodes):
            plaunch = subprocess.Popen(ec2launch)
            plaunch.wait()

        num_nodes = num_nodes*2

    # stop all instances and EBS volumes
    print "stopping all instances"
    inst_ids = [id[0] for id in current_ids]
    term = 'ec2kill ' + ' '.join(inst_ids)
    term = term.split()
    p = subprocess.Popen(term)
    p.wait()
    outfile.close()

def run_test(pair, queue):
    sshcon = ssh.Connection(host=pair[0][1], username='ubuntu',
                            private_key=PRIV_KEY)
    print "%s -> %s" % (pair[0][0], pair[1][0])
    output = sshcon.execute('iperf -c %s' % pair[1][2])
    result = ''.join(output[-1].split(' ')[-2:])
    result = result.strip('\n')
    row = [pair[0][0], pair[1][0], result]
    queue.put(row)
    sshcon.close()



def wait_for_launch():
    # wait until the instances are launched
    while True:
        instance_info = get_instance_info()
        # fail-safe
        assert len(instance_info) < MAX_NODES*2

        tokens = [t for t in ''.join(instance_info).split('\t') if t]
        if 'pending' not in tokens:
            time.sleep(10)
            return instance_info
        else:
            time.sleep(10)

def setup_instance(id):
    """Install dependencies and run iperf daemon an instance"""
    print "setting up %s" % id[0]
    while True:
        try:
            sshcon = ssh.Connection(host=id[1], username='ubuntu',
                                    private_key=PRIV_KEY)
        except:
            continue
        break
    print "installing iperf"
    sshcon.execute('sudo apt-get install -y --install-recommends ' \
                    'iperf')
    print "running iperf daemon"
    try:
        sshcon.execute('iperf -s -D', timeout=1.0)
    except:
        pass
    sshcon.close()


def get_instance_info():
    """Get instance info from ec2-describe-instances"""
    p = subprocess.Popen(ec2describe, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT)
    p.wait()
    return [line for line in p.stdout.readlines() if \
        line.startswith('INSTANCE') and ('running' in line.split('\t') or \
        'pending' in line.split('\t'))]

def get_instance_ids(info, exclude=None):
    """Get tuples of (instanceid, externalip, internalip) for each instance"""
    instance_ids = []
    for line in info:
        tokens = line.split('\t')
        id = (tokens[1], tokens[16], tokens[17])
        instance_ids.append((tokens[1], tokens[16], tokens[17]))
    return instance_ids

if __name__ == '__main__':
    main()
